{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NOTE:  THIS NOTEBOOK WILL TAKE A 30 MINUTES TO COMPLETE.\n",
    "\n",
    "# PLEASE BE PATIENT."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-Tuning a BERT Model and Create a Text Classifier\n",
    "\n",
    "In the previous section, we've already performed the Feature Engineering to create BERT embeddings from the `reviews_body` text using the pre-trained BERT model, and split the dataset into train, validation and test files. To optimize for Tensorflow training, we saved the files in TFRecord format. \n",
    "\n",
    "Now, let’s fine-tune the BERT model to our Customer Reviews Dataset and add a new classification layer to predict the `star_rating` for a given `review_body`.\n",
    "\n",
    "![BERT Training](img/bert_training.png)\n",
    "\n",
    "As mentioned earlier, BERT’s attention mechanism is called a Transformer. This is, not coincidentally, the name of the popular BERT Python library, “Transformers,” maintained by a company called HuggingFace. We will use a variant of BERT called [DistilBert](https://arxiv.org/pdf/1910.01108.pdf) which requires less memory and compute, but maintains very good accuracy on our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q sagemaker==2.20.0\n",
    "!pip install -q smdebug==1.0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import pandas as pd\n",
    "\n",
    "sess   = sagemaker.Session()\n",
    "bucket = sess.default_bucket()\n",
    "role = sagemaker.get_execution_role()\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "sm = boto3.Session().client(service_name='sagemaker', region_name=region)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# _PRE-REQUISITE: You need to have succesfully run the notebooks in the `PREPARE` section before you continue with this notebook._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r processed_train_data_s3_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    processed_train_data_s3_uri\n",
    "except NameError:\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')\n",
    "    print('[ERROR] Please run the notebooks in the PREPARE section before you continue.')\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-train\n"
     ]
    }
   ],
   "source": [
    "print(processed_train_data_s3_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r processed_validation_data_s3_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    processed_validation_data_s3_uri\n",
    "except NameError:\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')\n",
    "    print('[ERROR] Please run the notebooks in the PREPARE section before you continue.')\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-validation\n"
     ]
    }
   ],
   "source": [
    "print(processed_validation_data_s3_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r processed_test_data_s3_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    processed_test_data_s3_uri\n",
    "except NameError:\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')\n",
    "    print('[ERROR] Please run the notebooks in the PREPARE section before you continue.')\n",
    "    print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-test\n"
     ]
    }
   ],
   "source": [
    "print(processed_test_data_s3_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Track the `Experiment`\n",
    "We will track every step of this experiment throughout the `prepare`, `train`, `optimize`, and `deploy`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Concepts\n",
    "\n",
    "**Experiment**: A collection of related Trials.  Add Trials to an Experiment that you wish to compare together.\n",
    "\n",
    "**Trial**: A description of a multi-step machine learning workflow. Each step in the workflow is described by a Trial Component. There is no relationship between Trial Components such as ordering.\n",
    "\n",
    "**Trial Component**: A description of a single step in a machine learning workflow. For example data cleaning, feature extraction, model training, model evaluation, etc.\n",
    "\n",
    "**Tracker**: A logger of information about a single TrialComponent.\n",
    "\n",
    "<img src=\"img/sagemaker-experiments.png\" width=\"90%\" align=\"left\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the `Experiment`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Experiment name: Amazon-Customer-Reviews-BERT-Experiment-1609549145\n"
     ]
    }
   ],
   "source": [
    "# import time\n",
    "# from smexperiments.experiment import Experiment\n",
    "\n",
    "# timestamp = int(time.time())\n",
    "\n",
    "# experiment = Experiment.create(\n",
    "#                 experiment_name='Amazon-Customer-Reviews-BERT-Experiment-{}'.format(timestamp),\n",
    "#                 description='Amazon Customer Reviews BERT Experiment', \n",
    "#                 sagemaker_boto_client=sm)\n",
    "\n",
    "# experiment_name = experiment.experiment_name\n",
    "# print('Experiment name: {}'.format(experiment_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the `Trial`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trial name: trial-1609549145\n"
     ]
    }
   ],
   "source": [
    "# import time\n",
    "# from smexperiments.trial import Trial\n",
    "\n",
    "# timestamp = int(time.time())\n",
    "\n",
    "# trial = Trial.create(trial_name='trial-{}'.format(timestamp),\n",
    "#                      experiment_name=experiment_name,\n",
    "#                      sagemaker_boto_client=sm)\n",
    "\n",
    "# trial_name = trial.trial_name\n",
    "# print('Trial name: {}'.format(trial_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the `prepare` Trial Component and Tracker\n",
    "Note:  A Trial Component is actually created through a Tracker.  This is a bit confusing, we know."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepare trial component name TrialComponent-2021-01-02-005905-zwpm\n"
     ]
    }
   ],
   "source": [
    "# from smexperiments.tracker import Tracker\n",
    "\n",
    "# tracker_prepare = Tracker.create(display_name='prepare', \n",
    "#                                  sagemaker_boto_client=sm)\n",
    "\n",
    "# prepare_trial_component_name = tracker_prepare.trial_component.trial_component_name\n",
    "# print('Prepare trial component name {}'.format(prepare_trial_component_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attach the `prepare` Trial Component and Tracker as a Component to the Trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trial.add_trial_component(tracker_prepare.trial_component)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Log All Parameters Used During `prepare` Phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %store -r raw_input_data_s3_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/amazon-reviews-pds/tsv/\n"
     ]
    }
   ],
   "source": [
    "# print(raw_input_data_s3_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrialComponent(sagemaker_boto_client=<botocore.client.SageMaker object at 0x7f73efcb3b00>,trial_component_name='TrialComponent-2021-01-02-005905-zwpm',display_name='prepare',tags=None,trial_component_arn='arn:aws:sagemaker:us-east-1:835319576252:experiment-trial-component/trialcomponent-2021-01-02-005905-zwpm',response_metadata={'RequestId': '9adde967-b6b1-4492-b1ce-205ebe1af314', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amzn-requestid': '9adde967-b6b1-4492-b1ce-205ebe1af314', 'content-type': 'application/x-amz-json-1.1', 'content-length': '129', 'date': 'Sat, 02 Jan 2021 00:59:06 GMT'}, 'RetryAttempts': 0},parameters={},input_artifacts={'raw_data_s3_uri': TrialComponentArtifact(value='s3://sagemaker-us-east-1-835319576252/amazon-reviews-pds/tsv/',media_type='s3/uri')},output_artifacts={})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# tracker_prepare.log_input(name='raw_data_s3_uri', \n",
    "#                           media_type='s3/uri', \n",
    "#                           value=raw_input_data_s3_uri)\n",
    "\n",
    "# # must save after logging\n",
    "# tracker_prepare.trial_component.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no stored variable or alias train_split_percentage\n"
     ]
    }
   ],
   "source": [
    "# %store -r train_split_percentage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train_split_percentage' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-ab49f02a6dd6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_split_percentage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'train_split_percentage' is not defined"
     ]
    }
   ],
   "source": [
    "# print(train_split_percentage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %store -r validation_split_percentage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(validation_split_percentage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %store -r test_split_percentage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(test_split_percentage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r max_seq_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n"
     ]
    }
   ],
   "source": [
    "print(max_seq_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %store -r balance_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(balance_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tracker_prepare.log_parameters({\n",
    "#     'max_seq_length': max_seq_length,\n",
    "#     'train_split_percentage': train_split_percentage,\n",
    "#     'validation_split_percentage': validation_split_percentage,\n",
    "#     'test_split_percentage': test_split_percentage, \n",
    "#     'balance_dataset': str(balance_dataset)\n",
    "# })\n",
    "\n",
    "# # must save after logging\n",
    "# tracker_prepare.trial_component.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tracker_prepare.log_output(name='train_data_s3_uri', \n",
    "#                            media_type='s3/uri', \n",
    "#                            value=processed_train_data_s3_uri)\n",
    "\n",
    "# tracker_prepare.log_output(name='validation_data_s3_uri', \n",
    "#                            media_type='s3/uri', \n",
    "#                            value=processed_validation_data_s3_uri)\n",
    "\n",
    "# tracker_prepare.log_output(name='test_data_s3_uri', \n",
    "#                            media_type='s3/uri', \n",
    "#                            value=processed_test_data_s3_uri)\n",
    "\n",
    "# # must save after logging\n",
    "# tracker_prepare.trial_component.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify the Dataset in S3\n",
    "We are using the train, validation, and test splits created in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-train\n",
      "2020-12-31 08:49:55     353032 part-algo-1-amazon_reviews_us_Digital_Ebook_Purchase_v1_01.tfrecord\n",
      "2020-12-31 08:49:55      10844 part-algo-1-amazon_reviews_us_Digital_Software_v1_00.tfrecord\n",
      "2020-12-31 08:49:55      11611 part-algo-1-amazon_reviews_us_Digital_Video_Games_v1_00.tfrecord\n"
     ]
    }
   ],
   "source": [
    "print(processed_train_data_s3_uri)\n",
    "\n",
    "!aws s3 ls $processed_train_data_s3_uri/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-validation\n",
      "2020-12-31 08:49:55      19982 part-algo-1-amazon_reviews_us_Digital_Ebook_Purchase_v1_01.tfrecord\n",
      "2020-12-31 08:49:55        646 part-algo-1-amazon_reviews_us_Digital_Software_v1_00.tfrecord\n",
      "2020-12-31 08:49:55        716 part-algo-1-amazon_reviews_us_Digital_Video_Games_v1_00.tfrecord\n"
     ]
    }
   ],
   "source": [
    "print(processed_validation_data_s3_uri)\n",
    "\n",
    "!aws s3 ls $processed_validation_data_s3_uri/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-test\n",
      "2020-12-31 08:49:56      20020 part-algo-1-amazon_reviews_us_Digital_Ebook_Purchase_v1_01.tfrecord\n",
      "2020-12-31 08:49:56        692 part-algo-1-amazon_reviews_us_Digital_Software_v1_00.tfrecord\n",
      "2020-12-31 08:49:56        678 part-algo-1-amazon_reviews_us_Digital_Video_Games_v1_00.tfrecord\n"
     ]
    }
   ],
   "source": [
    "print(processed_test_data_s3_uri)\n",
    "\n",
    "!aws s3 ls $processed_test_data_s3_uri/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify S3 `Distribution Strategy`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-train', 'S3DataDistributionType': 'ShardedByS3Key'}}}\n",
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-validation', 'S3DataDistributionType': 'ShardedByS3Key'}}}\n",
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-12-31-08-43-25-560/output/bert-test', 'S3DataDistributionType': 'ShardedByS3Key'}}}\n"
     ]
    }
   ],
   "source": [
    "from sagemaker.inputs import TrainingInput\n",
    "\n",
    "s3_input_train_data = TrainingInput(s3_data=processed_train_data_s3_uri, \n",
    "                                         distribution='ShardedByS3Key') \n",
    "s3_input_validation_data = TrainingInput(s3_data=processed_validation_data_s3_uri, \n",
    "                                              distribution='ShardedByS3Key')\n",
    "s3_input_test_data = TrainingInput(s3_data=processed_test_data_s3_uri, \n",
    "                                        distribution='ShardedByS3Key')\n",
    "\n",
    "print(s3_input_train_data.config)\n",
    "print(s3_input_validation_data.config)\n",
    "print(s3_input_test_data.config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Hyper-Parameters for Classification Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n"
     ]
    }
   ],
   "source": [
    "print(max_seq_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs=3\n",
    "learning_rate=0.00001\n",
    "epsilon=0.00000001\n",
    "train_batch_size=128\n",
    "validation_batch_size=128\n",
    "test_batch_size=128\n",
    "train_steps_per_epoch=100\n",
    "validation_steps=100\n",
    "test_steps=100\n",
    "train_instance_count=1\n",
    "train_instance_type='ml.p3.2xlarge'\n",
    "train_volume_size=1024\n",
    "use_xla=True\n",
    "use_amp=True\n",
    "freeze_bert_layer=False\n",
    "enable_sagemaker_debugger=True\n",
    "enable_checkpointing=False\n",
    "enable_tensorboard=False\n",
    "input_mode='File'\n",
    "run_validation=True\n",
    "run_test=True\n",
    "run_sample_predictions=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Metrics To Track Model Performance\n",
    "\n",
    "These sample log lines...\n",
    "```\n",
    "45/50 [=====>..] - ETA: 3s - loss: 0.425 - accuracy: 0.881\n",
    "50/50 [=======>] - ETA: 0s - val_loss: 0.407 - val_accuracy: 0.885\n",
    "```\n",
    "...will produce the following 4 metrics in CloudWatch:\n",
    "\n",
    "`loss` = 0.425\n",
    "\n",
    "`accuracy` = 0.881\n",
    "\n",
    "`val_loss` = 0.407\n",
    "\n",
    "`val_accuracy` = 0.885"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/cloudwatch_train_accuracy.png\" width=\"50%\" align=\"left\">\n",
    "\n",
    "<img src=\"img/cloudwatch_train_loss.png\" width=\"50%\" align=\"left\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_definitions = [\n",
    "     {'Name': 'train:loss', 'Regex': 'loss: ([0-9\\\\.]+)'},\n",
    "     {'Name': 'train:accuracy', 'Regex': 'accuracy: ([0-9\\\\.]+)'},\n",
    "     {'Name': 'validation:loss', 'Regex': 'val_loss: ([0-9\\\\.]+)'},\n",
    "     {'Name': 'validation:accuracy', 'Regex': 'val_accuracy: ([0-9\\\\.]+)'},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup SageMaker Debugger\n",
    "Define Debugger Rules as deccribed here:  https://docs.aws.amazon.com/sagemaker/latest/dg/debugger-built-in-rules.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from sagemaker.debugger import Rule\n",
    "# from sagemaker.debugger import rule_configs\n",
    "# from sagemaker.debugger import CollectionConfig\n",
    "from sagemaker.debugger import DebuggerHookConfig\n",
    "\n",
    "# rules=[\n",
    "#         Rule.sagemaker(\n",
    "#             rule_configs.loss_not_decreasing(),\n",
    "#             rule_parameters={\n",
    "#                 'collection_names': 'losses,metrics',\n",
    "#                 'use_losses_collection': 'true',\n",
    "#                 'num_steps': '10',\n",
    "#                 'diff_percent': '50'\n",
    "#             },\n",
    "#             collections_to_save=[\n",
    "#                 CollectionConfig(name='losses',\n",
    "#                                  parameters={\n",
    "#                                      'save_interval': '10',\n",
    "#                                  }),\n",
    "#                 CollectionConfig(name='metrics',\n",
    "#                                  parameters={\n",
    "#                                      'save_interval': '10',\n",
    "#                                  })\n",
    "#             ]\n",
    "#         ),\n",
    "#         Rule.sagemaker(\n",
    "#             rule_configs.overtraining(),\n",
    "#             rule_parameters={\n",
    "#                 'collection_names': 'losses,metrics',\n",
    "#                 'patience_train': '10',\n",
    "#                 'patience_validation': '10',\n",
    "#                 'delta': '0.5'\n",
    "#             },\n",
    "#             collections_to_save=[\n",
    "#                 CollectionConfig(name='losses',\n",
    "#                                  parameters={\n",
    "#                                      'save_interval': '10',\n",
    "#                                  }),\n",
    "#                 CollectionConfig(name='metrics',\n",
    "#                                  parameters={\n",
    "#                                      'save_interval': '10',\n",
    "#                                  })\n",
    "#             ]\n",
    "#         )\n",
    "#     ]\n",
    "\n",
    "hook_config = DebuggerHookConfig(\n",
    "    hook_parameters={\n",
    "        'save_interval': '10', # number of steps\n",
    "        'export_tensorboard': 'true',\n",
    "        'tensorboard_dir': 'hook_tensorboard/',\n",
    "    })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Debugger Rules\n",
    "We specify the following rules:\n",
    "\n",
    "* loss_not_decreasing: checks if loss is decreasing and triggers if the loss has not decreased by a certain persentage in the last few iterations\n",
    "* LowGPUUtilization: checks if GPU is under-utilizated\n",
    "* ProfilerReport: runs the entire set of performance rules and create a final output report with further insights and recommendations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import Rule, ProfilerRule, rule_configs\n",
    "\n",
    "rules=[ \n",
    "    Rule.sagemaker(rule_configs.loss_not_decreasing()),\n",
    "    Rule.sagemaker(rule_configs.overtraining()),\n",
    "    ProfilerRule.sagemaker(rule_configs.ProfilerReport()),\n",
    "    ProfilerRule.sagemaker(rule_configs.LowGPUUtilization())\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify a Debugger profiler configuration\n",
    "\n",
    "The following configuration will capture system metrics at 500 milliseconds. The system metrics include utilization per CPU, GPU, memory utilization per CPU, GPU as well I/O and network.\n",
    "\n",
    "Debugger will capture detailed profiling information from step 5 to step 15. This information includes Horovod metrics, dataloading, preprocessing, operators running on CPU and GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import ProfilerConfig, FrameworkProfile\n",
    "\n",
    "profiler_config = ProfilerConfig(\n",
    "    system_monitor_interval_millis=500,\n",
    "    framework_profile_params=FrameworkProfile(local_path=\"/opt/ml/output/profiler/\", start_step=5, num_steps=10)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify Checkpoint S3 Location\n",
    "This is used for Spot Instances Training.  If nodes are replaced, the new node will start training from the latest checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-835319576252/checkpoints/6be259c8-0b8a-4f75-998d-b8c51038898c/\n"
     ]
    }
   ],
   "source": [
    "import uuid\n",
    "\n",
    "checkpoint_s3_prefix = 'checkpoints/{}'.format(str(uuid.uuid4()))\n",
    "checkpoint_s3_uri = 's3://{}/{}/'.format(bucket, checkpoint_s3_prefix)\n",
    "\n",
    "print(checkpoint_s3_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Our BERT + TensorFlow Script to Run on SageMaker\n",
    "Prepare our TensorFlow model to run on the managed SageMaker service"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtime\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mrandom\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mpandas\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mpd\u001b[39;49;00m\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mglob\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m glob\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mpprint\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36margparse\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mjson\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msubprocess\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msys\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mos\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtensorflow\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mtf\u001b[39;49;00m\r\n",
      "\u001b[37m#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'tensorflow==2.1.0'])\u001b[39;49;00m\r\n",
      "subprocess.check_call([sys.executable, \u001b[33m'\u001b[39;49;00m\u001b[33m-m\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mpip\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33minstall\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mtransformers==2.8.0\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "\u001b[37m#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sagemaker-tensorflow==2.3.0.1.6.1'])\u001b[39;49;00m\r\n",
      "\u001b[37m#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sagemaker-tensorflow==2.1.0.1.0.0'])\u001b[39;49;00m\r\n",
      "\u001b[37m#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'smdebug==0.9.3'])\u001b[39;49;00m\r\n",
      "subprocess.check_call([sys.executable, \u001b[33m'\u001b[39;49;00m\u001b[33m-m\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mpip\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33minstall\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33msmdebug==1.0.1\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "subprocess.check_call([sys.executable, \u001b[33m'\u001b[39;49;00m\u001b[33m-m\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mpip\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33minstall\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mscikit-learn==0.23.1\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "subprocess.check_call([sys.executable, \u001b[33m'\u001b[39;49;00m\u001b[33m-m\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mpip\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33minstall\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mmatplotlib==3.2.1\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtransformers\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m DistilBertTokenizer\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtransformers\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m TFDistilBertForSequenceClassification\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtransformers\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m TextClassificationPipeline\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtransformers\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mconfiguration_distilbert\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m DistilBertConfig\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtensorflow\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mkeras\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mcallbacks\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m ModelCheckpoint\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtensorflow\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mkeras\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mmodels\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m load_model\r\n",
      "\u001b[37m#from tensorflow.keras.mixed_precision import experimental as mixed_precision\u001b[39;49;00m\r\n",
      "\r\n",
      "\r\n",
      "CLASSES = [\u001b[34m1\u001b[39;49;00m, \u001b[34m2\u001b[39;49;00m, \u001b[34m3\u001b[39;49;00m, \u001b[34m4\u001b[39;49;00m, \u001b[34m5\u001b[39;49;00m]\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mselect_data_and_label_from_record\u001b[39;49;00m(record):\r\n",
      "    x = {\r\n",
      "        \u001b[33m'\u001b[39;49;00m\u001b[33minput_ids\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m: record[\u001b[33m'\u001b[39;49;00m\u001b[33minput_ids\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m],\r\n",
      "        \u001b[33m'\u001b[39;49;00m\u001b[33minput_mask\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m: record[\u001b[33m'\u001b[39;49;00m\u001b[33minput_mask\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m],\r\n",
      "        \u001b[33m'\u001b[39;49;00m\u001b[33msegment_ids\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m: record[\u001b[33m'\u001b[39;49;00m\u001b[33msegment_ids\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "    }\r\n",
      "\r\n",
      "    y = record[\u001b[33m'\u001b[39;49;00m\u001b[33mlabel_ids\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "\r\n",
      "    \u001b[34mreturn\u001b[39;49;00m (x, y)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mfile_based_input_dataset_builder\u001b[39;49;00m(channel,\r\n",
      "                                     input_filenames,\r\n",
      "                                     pipe_mode,\r\n",
      "                                     is_training,\r\n",
      "                                     drop_remainder,\r\n",
      "                                     batch_size,\r\n",
      "                                     epochs,\r\n",
      "                                     steps_per_epoch,\r\n",
      "                                     max_seq_length):\r\n",
      "\r\n",
      "    \u001b[37m# For training, we want a lot of parallel reading and shuffling.\u001b[39;49;00m\r\n",
      "    \u001b[37m# For eval, we want no shuffling and parallel reading doesn't matter.\u001b[39;49;00m\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m pipe_mode:\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m***** Using pipe_mode with channel \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(channel))\r\n",
      "        \u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36msagemaker_tensorflow\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m PipeModeDataset\r\n",
      "        dataset = PipeModeDataset(channel=channel,\r\n",
      "                                  record_format=\u001b[33m'\u001b[39;49;00m\u001b[33mTFRecord\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    \u001b[34melse\u001b[39;49;00m:\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m***** Using input_filenames \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(input_filenames))\r\n",
      "        dataset = tf.data.TFRecordDataset(input_filenames)\r\n",
      "\r\n",
      "    dataset = dataset.repeat(epochs * steps_per_epoch * \u001b[34m100\u001b[39;49;00m)\r\n",
      "\u001b[37m#    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\u001b[39;49;00m\r\n",
      "\r\n",
      "    name_to_features = {\r\n",
      "      \u001b[33m\"\u001b[39;49;00m\u001b[33minput_ids\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: tf.io.FixedLenFeature([max_seq_length], tf.int64),\r\n",
      "      \u001b[33m\"\u001b[39;49;00m\u001b[33minput_mask\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: tf.io.FixedLenFeature([max_seq_length], tf.int64),\r\n",
      "      \u001b[33m\"\u001b[39;49;00m\u001b[33msegment_ids\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: tf.io.FixedLenFeature([max_seq_length], tf.int64),\r\n",
      "      \u001b[33m\"\u001b[39;49;00m\u001b[33mlabel_ids\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: tf.io.FixedLenFeature([], tf.int64),\r\n",
      "    }\r\n",
      "\r\n",
      "    \u001b[34mdef\u001b[39;49;00m \u001b[32m_decode_record\u001b[39;49;00m(record, name_to_features):\r\n",
      "        \u001b[33m\"\"\"Decodes a record to a TensorFlow example.\"\"\"\u001b[39;49;00m\r\n",
      "        record = tf.io.parse_single_example(record, name_to_features)\r\n",
      "        \u001b[37m# TODO:  wip/bert/bert_attention_head_view/train.py\u001b[39;49;00m\r\n",
      "        \u001b[37m# Convert input_ids into input_tokens with DistilBert vocabulary \u001b[39;49;00m\r\n",
      "        \u001b[37m#  if hook.get_collections()['all'].save_config.should_save_step(modes.EVAL, hook.mode_steps[modes.EVAL]):\u001b[39;49;00m\r\n",
      "        \u001b[37m#    hook._write_raw_tensor_simple(\"input_tokens\", input_tokens)\u001b[39;49;00m\r\n",
      "        \u001b[34mreturn\u001b[39;49;00m record\r\n",
      "    \r\n",
      "    dataset = dataset.apply(\r\n",
      "        tf.data.experimental.map_and_batch(\r\n",
      "          \u001b[34mlambda\u001b[39;49;00m record: _decode_record(record, name_to_features),\r\n",
      "          batch_size=batch_size,\r\n",
      "          drop_remainder=drop_remainder,\r\n",
      "          num_parallel_calls=tf.data.experimental.AUTOTUNE))\r\n",
      "\r\n",
      "\u001b[37m#    dataset.cache()\u001b[39;49;00m\r\n",
      "\r\n",
      "    dataset = dataset.shuffle(buffer_size=\u001b[34m1000\u001b[39;49;00m,\r\n",
      "                              reshuffle_each_iteration=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "\r\n",
      "    row_count = \u001b[34m0\u001b[39;49;00m\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m**************** \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m *****************\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(channel))\r\n",
      "    \u001b[34mfor\u001b[39;49;00m row \u001b[35min\u001b[39;49;00m dataset.as_numpy_iterator():\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(row)\r\n",
      "        \u001b[34mif\u001b[39;49;00m row_count == \u001b[34m5\u001b[39;49;00m:\r\n",
      "            \u001b[34mbreak\u001b[39;49;00m\r\n",
      "        row_count = row_count + \u001b[34m1\u001b[39;49;00m\r\n",
      "\r\n",
      "    \u001b[34mreturn\u001b[39;49;00m dataset\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mload_checkpoint_model\u001b[39;49;00m(checkpoint_path):\r\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mglob\u001b[39;49;00m\r\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mos\u001b[39;49;00m\r\n",
      "    \r\n",
      "    glob_pattern = os.path.join(checkpoint_path, \u001b[33m'\u001b[39;49;00m\u001b[33m*.h5\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mglob pattern \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(glob_pattern))\r\n",
      "\r\n",
      "    list_of_checkpoint_files = glob.glob(glob_pattern)\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mList of checkpoint files \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(list_of_checkpoint_files))\r\n",
      "    \r\n",
      "    latest_checkpoint_file = \u001b[36mmax\u001b[39;49;00m(list_of_checkpoint_files)\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mLatest checkpoint file \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(latest_checkpoint_file))\r\n",
      "\r\n",
      "    initial_epoch_number_str = latest_checkpoint_file.rsplit(\u001b[33m'\u001b[39;49;00m\u001b[33m_\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[34m1\u001b[39;49;00m)[-\u001b[34m1\u001b[39;49;00m].split(\u001b[33m'\u001b[39;49;00m\u001b[33m.h5\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)[\u001b[34m0\u001b[39;49;00m]\r\n",
      "    initial_epoch_number = \u001b[36mint\u001b[39;49;00m(initial_epoch_number_str)\r\n",
      "\r\n",
      "    loaded_model = TFDistilBertForSequenceClassification.from_pretrained(\r\n",
      "                                               latest_checkpoint_file,\r\n",
      "                                               config=config)\r\n",
      "\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mloaded_model \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(loaded_model))\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33minitial_epoch_number \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(initial_epoch_number))\r\n",
      "    \r\n",
      "    \u001b[34mreturn\u001b[39;49;00m loaded_model, initial_epoch_number\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mif\u001b[39;49;00m \u001b[31m__name__\u001b[39;49;00m == \u001b[33m'\u001b[39;49;00m\u001b[33m__main__\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m:\r\n",
      "    parser = argparse.ArgumentParser()\r\n",
      "\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--train_data\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, \r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CHANNEL_TRAIN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--validation_data\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, \r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CHANNEL_VALIDATION\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--test_data\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m,\r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CHANNEL_TEST\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--output_dir\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m,\r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_OUTPUT_DIR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--hosts\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mlist\u001b[39;49;00m, \r\n",
      "                        default=json.loads(os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_HOSTS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]))\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--current_host\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, \r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CURRENT_HOST\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])    \r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--num_gpus\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, \r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_NUM_GPUS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--checkpoint_base_path\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, \r\n",
      "                        default=\u001b[33m'\u001b[39;49;00m\u001b[33m/opt/ml/checkpoints\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--use_xla\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--use_amp\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--max_seq_length\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m64\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--train_batch_size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m128\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--validation_batch_size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m256\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--test_batch_size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m256\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--epochs\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m2\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--learning_rate\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m0.00003\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--epsilon\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34m0.00000001\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--train_steps_per_epoch\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mNone\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--validation_steps\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mNone\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--test_steps\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mNone\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--freeze_bert_layer\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--enable_sagemaker_debugger\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--run_validation\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)    \r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--run_test\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)    \r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--run_sample_predictions\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--enable_tensorboard\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)        \r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--enable_checkpointing\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36meval\u001b[39;49;00m,\r\n",
      "                        default=\u001b[34mFalse\u001b[39;49;00m)    \r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--output_data_dir\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[37m# This is unused\u001b[39;49;00m\r\n",
      "                        \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m,\r\n",
      "                        default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_OUTPUT_DATA_DIR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    \r\n",
      "    \u001b[37m# This points to the S3 location - this should not be used by our code\u001b[39;49;00m\r\n",
      "    \u001b[37m# We should use /opt/ml/model/ instead\u001b[39;49;00m\r\n",
      "    \u001b[37m# parser.add_argument('--model_dir', \u001b[39;49;00m\r\n",
      "    \u001b[37m#                     type=str, \u001b[39;49;00m\r\n",
      "    \u001b[37m#                     default=os.environ['SM_MODEL_DIR'])\u001b[39;49;00m\r\n",
      "     \r\n",
      "    args, _ = parser.parse_known_args()\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\u001b[39;49;00m\u001b[33mArgs:\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m) \r\n",
      "    \u001b[36mprint\u001b[39;49;00m(args)\r\n",
      "    \r\n",
      "    env_var = os.environ \r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\u001b[39;49;00m\u001b[33mEnvironment Variables:\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m) \r\n",
      "    pprint.pprint(\u001b[36mdict\u001b[39;49;00m(env_var), width = \u001b[34m1\u001b[39;49;00m) \r\n",
      "\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mSM_TRAINING_ENV \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(env_var[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_TRAINING_ENV\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]))\r\n",
      "    sm_training_env_json = json.loads(env_var[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_TRAINING_ENV\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    is_master = sm_training_env_json[\u001b[33m'\u001b[39;49;00m\u001b[33mis_master\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mis_master \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(is_master))\r\n",
      "    \r\n",
      "    train_data = args.train_data\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtrain_data \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(train_data))\r\n",
      "    validation_data = args.validation_data\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mvalidation_data \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(validation_data))\r\n",
      "    test_data = args.test_data\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtest_data \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(test_data))    \r\n",
      "    local_model_dir = os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_MODEL_DIR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "    output_dir = args.output_dir\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33moutput_dir \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(output_dir))    \r\n",
      "    hosts = args.hosts\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mhosts \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(hosts))    \r\n",
      "    current_host = args.current_host\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mcurrent_host \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(current_host))    \r\n",
      "    num_gpus = args.num_gpus\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mnum_gpus \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(num_gpus))\r\n",
      "    job_name = os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSAGEMAKER_JOB_NAME\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mjob_name \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(job_name))    \r\n",
      "    use_xla = args.use_xla\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33muse_xla \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(use_xla))    \r\n",
      "    use_amp = args.use_amp\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33muse_amp \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(use_amp))    \r\n",
      "    max_seq_length = args.max_seq_length\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mmax_seq_length \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(max_seq_length))    \r\n",
      "    train_batch_size = args.train_batch_size\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtrain_batch_size \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(train_batch_size))    \r\n",
      "    validation_batch_size = args.validation_batch_size\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mvalidation_batch_size \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(validation_batch_size))    \r\n",
      "    test_batch_size = args.test_batch_size\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtest_batch_size \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(test_batch_size))    \r\n",
      "    epochs = args.epochs\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mepochs \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(epochs))    \r\n",
      "    learning_rate = args.learning_rate\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mlearning_rate \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(learning_rate))    \r\n",
      "    epsilon = args.epsilon\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mepsilon \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(epsilon))    \r\n",
      "    train_steps_per_epoch = args.train_steps_per_epoch\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtrain_steps_per_epoch \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(train_steps_per_epoch))    \r\n",
      "    validation_steps = args.validation_steps\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mvalidation_steps \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(validation_steps))    \r\n",
      "    test_steps = args.test_steps\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtest_steps \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(test_steps))    \r\n",
      "    freeze_bert_layer = args.freeze_bert_layer\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mfreeze_bert_layer \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(freeze_bert_layer))    \r\n",
      "    enable_sagemaker_debugger = args.enable_sagemaker_debugger\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33menable_sagemaker_debugger \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(enable_sagemaker_debugger))    \r\n",
      "    run_validation = args.run_validation\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mrun_validation \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(run_validation))    \r\n",
      "    run_test = args.run_test\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mrun_test \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(run_test))    \r\n",
      "    run_sample_predictions = args.run_sample_predictions\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mrun_sample_predictions \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(run_sample_predictions))\r\n",
      "    enable_tensorboard = args.enable_tensorboard\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33menable_tensorboard \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(enable_tensorboard))       \r\n",
      "    enable_checkpointing = args.enable_checkpointing\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33menable_checkpointing \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(enable_checkpointing))    \r\n",
      "\r\n",
      "    checkpoint_base_path = args.checkpoint_base_path\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mcheckpoint_base_path \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(checkpoint_base_path))\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m is_master:\r\n",
      "        checkpoint_path = checkpoint_base_path\r\n",
      "    \u001b[34melse\u001b[39;49;00m:\r\n",
      "        checkpoint_path = \u001b[33m'\u001b[39;49;00m\u001b[33m/tmp/checkpoints\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m        \r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mcheckpoint_path \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(checkpoint_path))\r\n",
      "    \r\n",
      "    \u001b[37m# Determine if PipeMode is enabled \u001b[39;49;00m\r\n",
      "    pipe_mode_str = os.environ.get(\u001b[33m'\u001b[39;49;00m\u001b[33mSM_INPUT_DATA_CONFIG\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    pipe_mode = (pipe_mode_str.find(\u001b[33m'\u001b[39;49;00m\u001b[33mPipe\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m) >= \u001b[34m0\u001b[39;49;00m)\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mUsing pipe_mode: \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(pipe_mode))\r\n",
      " \r\n",
      "    \u001b[37m# Model Output \u001b[39;49;00m\r\n",
      "    transformer_fine_tuned_model_path = os.path.join(local_model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mtransformers/fine-tuned/\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    os.makedirs(transformer_fine_tuned_model_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[37m# SavedModel Output\u001b[39;49;00m\r\n",
      "    tensorflow_saved_model_path = os.path.join(local_model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mtensorflow/saved_model/0\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    os.makedirs(tensorflow_saved_model_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[37m# Tensorboard Logs \u001b[39;49;00m\r\n",
      "    tensorboard_logs_path = os.path.join(local_model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mtensorboard/\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    os.makedirs(tensorboard_logs_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[37m# Commented out due to incompatibility with transformers library (possibly)\u001b[39;49;00m\r\n",
      "    \u001b[37m# Set the global precision mixed_precision policy to \"mixed_float16\"    \u001b[39;49;00m\r\n",
      "\u001b[37m#    mixed_precision_policy = 'mixed_float16'\u001b[39;49;00m\r\n",
      "\u001b[37m#    print('Mixed precision policy {}'.format(mixed_precision_policy))\u001b[39;49;00m\r\n",
      "\u001b[37m#    policy = mixed_precision.Policy(mixed_precision_policy)\u001b[39;49;00m\r\n",
      "\u001b[37m#    mixed_precision.set_policy(policy)    \u001b[39;49;00m\r\n",
      "    \r\n",
      "    distributed_strategy = tf.distribute.MirroredStrategy()\r\n",
      "    \u001b[37m# Comment out when using smdebug as smdebug does not support MultiWorkerMirroredStrategy() as of smdebug 0.8.0\u001b[39;49;00m\r\n",
      "    \u001b[37m#distributed_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\u001b[39;49;00m\r\n",
      "    \u001b[34mwith\u001b[39;49;00m distributed_strategy.scope():\r\n",
      "        tf.config.optimizer.set_jit(use_xla)\r\n",
      "        tf.config.optimizer.set_experimental_options({\u001b[33m\"\u001b[39;49;00m\u001b[33mauto_mixed_precision\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: use_amp})\r\n",
      "\r\n",
      "        train_data_filenames = glob(os.path.join(train_data, \u001b[33m'\u001b[39;49;00m\u001b[33m*.tfrecord\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m))\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtrain_data_filenames \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(train_data_filenames))\r\n",
      "        train_dataset = file_based_input_dataset_builder(\r\n",
      "            channel=\u001b[33m'\u001b[39;49;00m\u001b[33mtrain\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "            input_filenames=train_data_filenames,\r\n",
      "            pipe_mode=pipe_mode,\r\n",
      "            is_training=\u001b[34mTrue\u001b[39;49;00m,\r\n",
      "            drop_remainder=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "            batch_size=train_batch_size,\r\n",
      "            epochs=epochs,\r\n",
      "            steps_per_epoch=train_steps_per_epoch,\r\n",
      "            max_seq_length=max_seq_length).map(select_data_and_label_from_record)\r\n",
      "\r\n",
      "        tokenizer = \u001b[34mNone\u001b[39;49;00m\r\n",
      "        config = \u001b[34mNone\u001b[39;49;00m\r\n",
      "        model = \u001b[34mNone\u001b[39;49;00m\r\n",
      "\r\n",
      "        \u001b[37m# This is required when launching many instances at once...  the urllib request seems to get denied periodically\u001b[39;49;00m\r\n",
      "        successful_download = \u001b[34mFalse\u001b[39;49;00m\r\n",
      "        retries = \u001b[34m0\u001b[39;49;00m\r\n",
      "        \u001b[34mwhile\u001b[39;49;00m (retries < \u001b[34m5\u001b[39;49;00m \u001b[35mand\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m successful_download):\r\n",
      "            \u001b[34mtry\u001b[39;49;00m:\r\n",
      "                tokenizer = DistilBertTokenizer.from_pretrained(\u001b[33m'\u001b[39;49;00m\u001b[33mdistilbert-base-uncased\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                config = DistilBertConfig.from_pretrained(\u001b[33m'\u001b[39;49;00m\u001b[33mdistilbert-base-uncased\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                                                          num_labels=\u001b[36mlen\u001b[39;49;00m(CLASSES))\r\n",
      "                model = TFDistilBertForSequenceClassification.from_pretrained(\u001b[33m'\u001b[39;49;00m\u001b[33mdistilbert-base-uncased\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                                                                              config=config)\r\n",
      "                successful_download = \u001b[34mTrue\u001b[39;49;00m\r\n",
      "                \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mSucessfully downloaded after \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m retries.\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(retries))\r\n",
      "            \u001b[34mexcept\u001b[39;49;00m:\r\n",
      "                retries = retries + \u001b[34m1\u001b[39;49;00m\r\n",
      "                random_sleep = random.randint(\u001b[34m1\u001b[39;49;00m, \u001b[34m30\u001b[39;49;00m)\r\n",
      "                \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mRetry #\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m.  Sleeping for \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m seconds\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(retries, random_sleep))\r\n",
      "                time.sleep(random_sleep)\r\n",
      "\r\n",
      "        callbacks = []\r\n",
      "\r\n",
      "        initial_epoch_number = \u001b[34m0\u001b[39;49;00m \r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m enable_checkpointing:\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m***** Checkpoint enabled *****\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            \r\n",
      "            os.makedirs(checkpoint_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)        \r\n",
      "            \u001b[34mif\u001b[39;49;00m os.listdir(checkpoint_path):\r\n",
      "                \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m***** Found checkpoint *****\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "                \u001b[36mprint\u001b[39;49;00m(checkpoint_path)\r\n",
      "                model, initial_epoch_number = load_checkpoint_model(checkpoint_path)\r\n",
      "                \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m***** Using checkpoint model \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m *****\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(model))\r\n",
      "                \r\n",
      "            checkpoint_callback = ModelCheckpoint(\r\n",
      "                    filepath=os.path.join(checkpoint_path, \u001b[33m'\u001b[39;49;00m\u001b[33mtf_model_\u001b[39;49;00m\u001b[33m{epoch:05d}\u001b[39;49;00m\u001b[33m.h5\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m),\r\n",
      "                    save_weights_only=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                    verbose=\u001b[34m1\u001b[39;49;00m,\r\n",
      "                    monitor=\u001b[33m'\u001b[39;49;00m\u001b[33mval_accuracy\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m*** CHECKPOINT CALLBACK \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m ***\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(checkpoint_callback))\r\n",
      "            callbacks.append(checkpoint_callback)\r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m tokenizer \u001b[35mor\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m model \u001b[35mor\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m config:\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mNot properly initialized...\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "        optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate, epsilon=epsilon)\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m** use_amp \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(use_amp))        \r\n",
      "        \u001b[34mif\u001b[39;49;00m use_amp:\r\n",
      "            \u001b[37m# loss scaling is currently required when using mixed precision\u001b[39;49;00m\r\n",
      "            optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, \u001b[33m'\u001b[39;49;00m\u001b[33mdynamic\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33menable_sagemaker_debugger \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(enable_sagemaker_debugger))\r\n",
      "        \u001b[34mif\u001b[39;49;00m enable_sagemaker_debugger:\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m*** DEBUGGING ***\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msmdebug\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mtensorflow\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36msmd\u001b[39;49;00m\r\n",
      "            \u001b[37m# This assumes that we specified debugger_hook_config\u001b[39;49;00m\r\n",
      "            debugger_callback = smd.KerasHook.create_from_json_file()\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m*** DEBUGGER CALLBACK \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m ***\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(debugger_callback))            \r\n",
      "            callbacks.append(debugger_callback)\r\n",
      "            optimizer = debugger_callback.wrap_optimizer(optimizer)\r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m enable_tensorboard:            \r\n",
      "            tensorboard_callback = tf.keras.callbacks.TensorBoard(\r\n",
      "                                                        log_dir=tensorboard_logs_path)\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m*** TENSORBOARD CALLBACK \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m ***\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(tensorboard_callback))\r\n",
      "            callbacks.append(tensorboard_callback)\r\n",
      "  \r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m*** OPTIMIZER \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m ***\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(optimizer))\r\n",
      "        \r\n",
      "        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "        metric = tf.keras.metrics.SparseCategoricalAccuracy(\u001b[33m'\u001b[39;49;00m\u001b[33maccuracy\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "        model.compile(optimizer=optimizer, loss=loss, metrics=[metric])\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mCompiled model \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(model))          \r\n",
      "        model.layers[\u001b[34m0\u001b[39;49;00m].trainable = \u001b[35mnot\u001b[39;49;00m freeze_bert_layer\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(model.summary())\r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m run_validation:\r\n",
      "            validation_data_filenames = glob(os.path.join(validation_data, \u001b[33m'\u001b[39;49;00m\u001b[33m*.tfrecord\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m))\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mvalidation_data_filenames \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(validation_data_filenames))\r\n",
      "            validation_dataset = file_based_input_dataset_builder(\r\n",
      "                channel=\u001b[33m'\u001b[39;49;00m\u001b[33mvalidation\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                input_filenames=validation_data_filenames,\r\n",
      "                pipe_mode=pipe_mode,\r\n",
      "                is_training=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                drop_remainder=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                batch_size=validation_batch_size,\r\n",
      "                epochs=epochs,\r\n",
      "                steps_per_epoch=validation_steps,\r\n",
      "                max_seq_length=max_seq_length).map(select_data_and_label_from_record)\r\n",
      "            \r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mStarting Training and Validation...\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            validation_dataset = validation_dataset.take(validation_steps)\r\n",
      "            train_and_validation_history = model.fit(train_dataset,\r\n",
      "                                                     shuffle=\u001b[34mTrue\u001b[39;49;00m,\r\n",
      "                                                     epochs=epochs,\r\n",
      "                                                     initial_epoch=initial_epoch_number,\r\n",
      "                                                     steps_per_epoch=train_steps_per_epoch,\r\n",
      "                                                     validation_data=validation_dataset,\r\n",
      "                                                     validation_steps=validation_steps,\r\n",
      "                                                     callbacks=callbacks)                                \r\n",
      "            \u001b[36mprint\u001b[39;49;00m(train_and_validation_history)\r\n",
      "        \u001b[34melse\u001b[39;49;00m: \u001b[37m# Not running validation\u001b[39;49;00m\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mStarting Training (Without Validation)...\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            train_history = model.fit(train_dataset,\r\n",
      "                                      shuffle=\u001b[34mTrue\u001b[39;49;00m,\r\n",
      "                                      epochs=epochs,\r\n",
      "                                      initial_epoch=initial_epoch_number,\r\n",
      "                                      steps_per_epoch=train_steps_per_epoch,\r\n",
      "                                      callbacks=callbacks)                \r\n",
      "            \u001b[36mprint\u001b[39;49;00m(train_history)\r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m run_test:\r\n",
      "            test_data_filenames = glob(os.path.join(test_data, \u001b[33m'\u001b[39;49;00m\u001b[33m*.tfrecord\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m))\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtest_data_filenames \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(test_data_filenames))\r\n",
      "            test_dataset = file_based_input_dataset_builder(\r\n",
      "                channel=\u001b[33m'\u001b[39;49;00m\u001b[33mtest\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                input_filenames=test_data_filenames,\r\n",
      "                pipe_mode=pipe_mode,\r\n",
      "                is_training=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                drop_remainder=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                batch_size=test_batch_size,\r\n",
      "                epochs=epochs,\r\n",
      "                steps_per_epoch=test_steps,\r\n",
      "                max_seq_length=max_seq_length).map(select_data_and_label_from_record)\r\n",
      "\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mStarting test...\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "            test_history = model.evaluate(test_dataset,\r\n",
      "                                          steps=test_steps,\r\n",
      "                                          callbacks=callbacks)\r\n",
      "                                 \r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mTest history \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(test_history))\r\n",
      "            \r\n",
      "        \u001b[37m# Save the Fine-Yuned Transformers Model as a New \"Pre-Trained\" Model\u001b[39;49;00m\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtransformer_fine_tuned_model_path \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(transformer_fine_tuned_model_path))   \r\n",
      "        model.save_pretrained(transformer_fine_tuned_model_path)\r\n",
      "\r\n",
      "        \u001b[37m# Save the TensorFlow SavedModel for Serving Predictions\u001b[39;49;00m\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mtensorflow_saved_model_path \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(tensorflow_saved_model_path))   \r\n",
      "        model.save(tensorflow_saved_model_path, save_format=\u001b[33m'\u001b[39;49;00m\u001b[33mtf\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "                \r\n",
      "        \u001b[37m# Copy inference.py and requirements.txt to the code/ directory\u001b[39;49;00m\r\n",
      "        \u001b[37m#   Note: This is required for the SageMaker Endpoint to pick them up.\u001b[39;49;00m\r\n",
      "        \u001b[37m#         This appears to be hard-coded and must be called code/\u001b[39;49;00m\r\n",
      "        inference_path = os.path.join(local_model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mcode/\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mCopying inference source files to \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(inference_path))\r\n",
      "        os.makedirs(inference_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)               \r\n",
      "        os.system(\u001b[33m'\u001b[39;49;00m\u001b[33mcp inference.py \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(inference_path))\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(glob(inference_path))        \r\n",
      "\u001b[37m#        os.system('cp requirements.txt {}/code'.format(inference_path))\u001b[39;49;00m\r\n",
      "        \r\n",
      "    \u001b[34mif\u001b[39;49;00m run_sample_predictions:\r\n",
      "        loaded_model = TFDistilBertForSequenceClassification.from_pretrained(transformer_fine_tuned_model_path,\r\n",
      "                                                                       id2label={\r\n",
      "                                                                        \u001b[34m0\u001b[39;49;00m: \u001b[34m1\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m1\u001b[39;49;00m: \u001b[34m2\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m2\u001b[39;49;00m: \u001b[34m3\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m3\u001b[39;49;00m: \u001b[34m4\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m4\u001b[39;49;00m: \u001b[34m5\u001b[39;49;00m\r\n",
      "                                                                       },\r\n",
      "                                                                       label2id={\r\n",
      "                                                                        \u001b[34m1\u001b[39;49;00m: \u001b[34m0\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m2\u001b[39;49;00m: \u001b[34m1\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m3\u001b[39;49;00m: \u001b[34m2\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m4\u001b[39;49;00m: \u001b[34m3\u001b[39;49;00m,\r\n",
      "                                                                        \u001b[34m5\u001b[39;49;00m: \u001b[34m4\u001b[39;49;00m\r\n",
      "                                                                       })\r\n",
      "\r\n",
      "        tokenizer = DistilBertTokenizer.from_pretrained(\u001b[33m'\u001b[39;49;00m\u001b[33mdistilbert-base-uncased\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "        \u001b[34mif\u001b[39;49;00m num_gpus >= \u001b[34m1\u001b[39;49;00m:\r\n",
      "            inference_device = \u001b[34m0\u001b[39;49;00m \u001b[37m# GPU 0\u001b[39;49;00m\r\n",
      "        \u001b[34melse\u001b[39;49;00m:\r\n",
      "            inference_device = -\u001b[34m1\u001b[39;49;00m \u001b[37m# CPU\u001b[39;49;00m\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33minference_device \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(inference_device))\r\n",
      "\r\n",
      "        inference_pipeline = TextClassificationPipeline(model=loaded_model, \r\n",
      "                                                        tokenizer=tokenizer,\r\n",
      "                                                        framework=\u001b[33m'\u001b[39;49;00m\u001b[33mtf\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                                                        device=inference_device)  \r\n",
      "\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mI loved it!  I will recommend this to everyone.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m, inference_pipeline(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mI loved it!  I will recommend this to everyone.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m))\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mIt\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\u001b[33ms OK.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m, inference_pipeline(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mIt\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\u001b[33ms OK.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m))\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mReally bad.  I hope they don\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\u001b[33mt make this anymore.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m, inference_pipeline(\u001b[33m\"\"\"\u001b[39;49;00m\u001b[33mReally bad.  I hope they don\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\u001b[33mt make this anymore.\u001b[39;49;00m\u001b[33m\"\"\"\u001b[39;49;00m))\r\n",
      "\r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mcsv\u001b[39;49;00m\r\n",
      "\r\n",
      "        df_test_reviews = pd.read_csv(\u001b[33m'\u001b[39;49;00m\u001b[33m./test_data/amazon_reviews_us_Digital_Software_v1_00.tsv.gz\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                                        delimiter=\u001b[33m'\u001b[39;49;00m\u001b[33m\\t\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \r\n",
      "                                        quoting=csv.QUOTE_NONE,\r\n",
      "                                        compression=\u001b[33m'\u001b[39;49;00m\u001b[33mgzip\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)[[\u001b[33m'\u001b[39;49;00m\u001b[33mreview_body\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mstar_rating\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]]\r\n",
      "\r\n",
      "        df_test_reviews = df_test_reviews.sample(n=\u001b[34m100\u001b[39;49;00m)\r\n",
      "        df_test_reviews.shape\r\n",
      "        df_test_reviews.head()\r\n",
      "        \r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mpandas\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mpd\u001b[39;49;00m\r\n",
      "\r\n",
      "        \u001b[34mdef\u001b[39;49;00m \u001b[32mpredict\u001b[39;49;00m(review_body):\r\n",
      "            prediction_map = inference_pipeline(review_body)\r\n",
      "            \u001b[34mreturn\u001b[39;49;00m prediction_map[\u001b[34m0\u001b[39;49;00m][\u001b[33m'\u001b[39;49;00m\u001b[33mlabel\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "\r\n",
      "        y_test = df_test_reviews[\u001b[33m'\u001b[39;49;00m\u001b[33mreview_body\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m].map(predict)\r\n",
      "        y_test\r\n",
      "        \r\n",
      "        y_actual = df_test_reviews[\u001b[33m'\u001b[39;49;00m\u001b[33mstar_rating\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]\r\n",
      "        y_actual\r\n",
      "\r\n",
      "        \u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36msklearn\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mmetrics\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m classification_report\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(classification_report(y_true=y_test, y_pred=y_actual))\r\n",
      "        \r\n",
      "        \u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36msklearn\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mmetrics\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m accuracy_score\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mAccuracy: \u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, accuracy_score(y_true=y_test, y_pred=y_actual))\r\n",
      "        \r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mmatplotlib\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mpyplot\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mplt\u001b[39;49;00m\r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mpandas\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mpd\u001b[39;49;00m\r\n",
      "\r\n",
      "        \u001b[34mdef\u001b[39;49;00m \u001b[32mplot_conf_mat\u001b[39;49;00m(cm, classes, title, cmap = plt.cm.Greens):\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(cm)\r\n",
      "            plt.imshow(cm, interpolation=\u001b[33m'\u001b[39;49;00m\u001b[33mnearest\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, cmap=cmap)\r\n",
      "            plt.title(title)\r\n",
      "            plt.colorbar()\r\n",
      "            tick_marks = np.arange(\u001b[36mlen\u001b[39;49;00m(classes))\r\n",
      "            plt.xticks(tick_marks, classes, rotation=\u001b[34m45\u001b[39;49;00m)\r\n",
      "            plt.yticks(tick_marks, classes)\r\n",
      "\r\n",
      "            fmt = \u001b[33m'\u001b[39;49;00m\u001b[33md\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\r\n",
      "            thresh = cm.max() / \u001b[34m2.\u001b[39;49;00m\r\n",
      "            \u001b[34mfor\u001b[39;49;00m i, j \u001b[35min\u001b[39;49;00m itertools.product(\u001b[36mrange\u001b[39;49;00m(cm.shape[\u001b[34m0\u001b[39;49;00m]), \u001b[36mrange\u001b[39;49;00m(cm.shape[\u001b[34m1\u001b[39;49;00m])):\r\n",
      "                plt.text(j, i, \u001b[36mformat\u001b[39;49;00m(cm[i, j], fmt),\r\n",
      "                horizontalalignment=\u001b[33m\"\u001b[39;49;00m\u001b[33mcenter\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m,\r\n",
      "                color=\u001b[33m\"\u001b[39;49;00m\u001b[33mblack\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m \u001b[34mif\u001b[39;49;00m cm[i, j] > thresh \u001b[34melse\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mblack\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "\r\n",
      "                plt.tight_layout()\r\n",
      "                plt.ylabel(\u001b[33m'\u001b[39;49;00m\u001b[33mTrue label\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "                plt.xlabel(\u001b[33m'\u001b[39;49;00m\u001b[33mPredicted label\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "                \r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mitertools\u001b[39;49;00m\r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mnumpy\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnp\u001b[39;49;00m\r\n",
      "        \u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36msklearn\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mmetrics\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m confusion_matrix\r\n",
      "        \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mmatplotlib\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mpyplot\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mplt\u001b[39;49;00m\r\n",
      "        \u001b[37m#%matplotlib inline\u001b[39;49;00m\r\n",
      "        \u001b[37m#%config InlineBackend.figure_format='retina'\u001b[39;49;00m\r\n",
      "\r\n",
      "        cm = confusion_matrix(y_true=y_test, y_pred=y_actual)\r\n",
      "\r\n",
      "        plt.figure()\r\n",
      "        fig, ax = plt.subplots(figsize=(\u001b[34m10\u001b[39;49;00m,\u001b[34m5\u001b[39;49;00m))\r\n",
      "        plot_conf_mat(cm, \r\n",
      "                      classes=[\u001b[33m'\u001b[39;49;00m\u001b[33m1\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33m2\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33m3\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33m4\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33m5\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m], \r\n",
      "                      title=\u001b[33m'\u001b[39;49;00m\u001b[33mConfusion Matrix\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "        \u001b[37m# Save the confusion matrix        \u001b[39;49;00m\r\n",
      "        plt.show()\r\n",
      "        \r\n",
      "        \u001b[37m# Model Output \u001b[39;49;00m\r\n",
      "        metrics_path = os.path.join(local_model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mmetrics/\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "        os.makedirs(metrics_path, exist_ok=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "        plt.savefig(\u001b[33m'\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m/confusion_matrix.png\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(metrics_path))\r\n"
     ]
    }
   ],
   "source": [
    "!pygmentize src-tf231-profiler/tf_bert_reviews.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37\n"
     ]
    }
   ],
   "source": [
    "tf_image = sagemaker.image_uris.retrieve('tensorflow', region=region, version='2.3.1', py_version='py37', image_scope='training', instance_type='ml.p3.2xlarge')\n",
    "print(tf_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arn:aws:iam::835319576252:role/service-role/AmazonSageMaker-ExecutionRole-20191006T135881\n"
     ]
    }
   ],
   "source": [
    "print(role)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "estimator = TensorFlow(entry_point='tf_bert_reviews.py',\n",
    "                       source_dir='src-tf231-profiler',\n",
    "                       role=role,\n",
    "                       image_uri=tf_image,\n",
    "                       instance_count=train_instance_count,\n",
    "                       instance_type=train_instance_type,\n",
    "                       volume_size=train_volume_size,\n",
    "#                        use_spot_instances=True,\n",
    "#                        max_wait=7200, # Seconds to wait for spot instances to become available\n",
    "                       checkpoint_s3_uri=checkpoint_s3_uri,\n",
    "#                       py_version='py37',\n",
    "#                       framework_version='2.3.1',\n",
    "                       hyperparameters={'epochs': epochs,\n",
    "                                        'learning_rate': learning_rate,\n",
    "                                        'epsilon': epsilon,\n",
    "                                        'train_batch_size': train_batch_size,\n",
    "                                        'validation_batch_size': validation_batch_size,\n",
    "                                        'test_batch_size': test_batch_size,                                             \n",
    "                                        'train_steps_per_epoch': train_steps_per_epoch,\n",
    "                                        'validation_steps': validation_steps,\n",
    "                                        'test_steps': test_steps,\n",
    "                                        'use_xla': use_xla,\n",
    "                                        'use_amp': use_amp,                                             \n",
    "                                        'max_seq_length': max_seq_length,\n",
    "                                        'freeze_bert_layer': freeze_bert_layer,\n",
    "                                        'enable_sagemaker_debugger': enable_sagemaker_debugger,\n",
    "                                        'enable_checkpointing': enable_checkpointing,\n",
    "                                        'enable_tensorboard': enable_tensorboard,                                        \n",
    "                                        'run_validation': run_validation,\n",
    "                                        'run_test': run_test,\n",
    "                                        'run_sample_predictions': run_sample_predictions},\n",
    "                       input_mode=input_mode,\n",
    "                       metric_definitions=metrics_definitions,\n",
    "                       rules=rules,\n",
    "                       debugger_hook_config=hook_config, \n",
    "                       profiler_config=profiler_config\n",
    "#                       max_run=7200, # number of seconds\n",
    "                      )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the `Experiment Config`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_config = {\n",
    "    'ExperimentName': experiment_name,\n",
    "    'TrialName': trial.trial_name,\n",
    "    'TrialComponentDisplayName': 'train'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Amazon-Customer-Reviews-BERT-Experiment-1609549145\n"
     ]
    }
   ],
   "source": [
    "print(experiment_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stored 'experiment_name' (str)\n"
     ]
    }
   ],
   "source": [
    "%store experiment_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trial-1609549145\n"
     ]
    }
   ],
   "source": [
    "print(trial_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stored 'trial_name' (str)\n"
     ]
    }
   ],
   "source": [
    "%store trial_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the Model on SageMaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:sagemaker.image_uris:Defaulting to the only supported framework/algorithm version: latest.\n",
      "INFO:sagemaker.image_uris:Ignoring unnecessary instance type: None.\n",
      "INFO:sagemaker.image_uris:Defaulting to the only supported framework/algorithm version: latest.\n",
      "INFO:sagemaker.image_uris:Ignoring unnecessary instance type: None.\n",
      "INFO:sagemaker.image_uris:Defaulting to the only supported framework/algorithm version: latest.\n",
      "INFO:sagemaker.image_uris:Ignoring unnecessary instance type: None.\n",
      "INFO:sagemaker.image_uris:Defaulting to the only supported framework/algorithm version: latest.\n",
      "INFO:sagemaker.image_uris:Ignoring unnecessary instance type: None.\n",
      "INFO:sagemaker:Creating training-job with name: tensorflow-training-2021-01-02-01-02-31-908\n"
     ]
    }
   ],
   "source": [
    "estimator.fit(inputs={'train': s3_input_train_data, \n",
    "                      'validation': s3_input_validation_data,\n",
    "                      'test': s3_input_test_data\n",
    "              },              \n",
    "              experiment_config=experiment_config,                   \n",
    "              wait=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Job Name:  tensorflow-training-2021-01-02-01-02-31-908\n"
     ]
    }
   ],
   "source": [
    "training_job_name = estimator.latest_training_job.name\n",
    "print('Training Job Name:  {}'.format(training_job_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region=us-east-1#/jobs/tensorflow-training-2021-01-02-01-02-31-908\">Training Job</a> After About 5 Minutes</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region={}#/jobs/{}\">Training Job</a> After About 5 Minutes</b>'.format(region, training_job_name)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a target=\"blank\" href=\"https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logStream:group=/aws/sagemaker/TrainingJobs;prefix=tensorflow-training-2021-01-02-01-02-31-908;streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a> After About 5 Minutes</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a target=\"blank\" href=\"https://console.aws.amazon.com/cloudwatch/home?region={}#logStream:group=/aws/sagemaker/TrainingJobs;prefix={};streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a> After About 5 Minutes</b>'.format(region, training_job_name)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a target=\"blank\" href=\"https://s3.console.aws.amazon.com/s3/buckets/sagemaker-us-east-1-835319576252/tensorflow-training-2021-01-02-01-02-31-908/?region=us-east-1&tab=overview\">S3 Output Data</a> After The Training Job Has Completed</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a target=\"blank\" href=\"https://s3.console.aws.amazon.com/s3/buckets/{}/{}/?region={}&tab=overview\">S3 Output Data</a> After The Training Job Has Completed</b>'.format(bucket, training_job_name, region)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a target=\"blank\" href=\"https://s3.console.aws.amazon.com/s3/buckets/sagemaker-us-east-1-835319576252/checkpoints/6be259c8-0b8a-4f75-998d-b8c51038898c/?region=us-east-1&tab=overview\">S3 Checkpoint Data</a> After The Training Job Has Completed</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a target=\"blank\" href=\"https://s3.console.aws.amazon.com/s3/buckets/{}/{}/?region={}&tab=overview\">S3 Checkpoint Data</a> After The Training Job Has Completed</b>'.format(bucket, checkpoint_s3_prefix, region)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2021-01-02 01:02:36 Starting - Launching requested ML instances.................\n",
      "2021-01-02 01:04:04 Starting - Preparing the instances for training.................\n",
      "2021-01-02 01:05:36 Downloading - Downloading input data........\n",
      "2021-01-02 01:06:22 Training - Downloading the training image.......\n",
      "2021-01-02 01:07:03 Training - Training image download completed. Training in progress..................................................................................\n",
      "2021-01-02 01:13:54 Uploading - Uploading generated training model...............\n",
      "2021-01-02 01:15:13 Completed - Training job completed\n",
      "CPU times: user 608 ms, sys: 21.7 ms, total: 630 ms\n",
      "Wall time: 12min 39s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "estimator.latest_training_job.wait(logs=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sage = boto3.client('sagemaker')\n",
    "# sage.describe_training_job(TrainingJobName='tensorflow-training-2020-12-18-18-01-49-949')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wait Until the ^^ Training Job ^^ Completes Above!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [INFO] _Feel free to continue to the next workshop section while this notebook is running._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "download: s3://sagemaker-us-east-1-835319576252/tensorflow-training-2021-01-02-01-02-31-908/output/model.tar.gz to ./model.tar.gz\n"
     ]
    }
   ],
   "source": [
    "!aws s3 cp s3://$bucket/$training_job_name/output/model.tar.gz ./model.tar.gz\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "metrics/\n",
      "metrics/confusion_matrix.png\n",
      "transformers/\n",
      "transformers/fine-tuned/\n",
      "transformers/fine-tuned/tf_model.h5\n",
      "transformers/fine-tuned/config.json\n",
      "tensorflow/\n",
      "tensorflow/saved_model/\n",
      "tensorflow/saved_model/0/\n",
      "tensorflow/saved_model/0/saved_model.pb\n",
      "tensorflow/saved_model/0/variables/\n",
      "tensorflow/saved_model/0/variables/variables.data-00000-of-00001\n",
      "tensorflow/saved_model/0/variables/variables.index\n",
      "tensorflow/saved_model/0/assets/\n",
      "tensorboard/\n",
      "code/\n",
      "code/inference.py\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p ./model/\n",
    "!tar -xvzf ./model.tar.gz -C ./model/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-02 01:16:10.094839: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1\n",
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['__saved_model_init_op']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['__saved_model_init_op'] tensor_info:\n",
      "        dtype: DT_INVALID\n",
      "        shape: unknown_rank\n",
      "        name: NoOp\n",
      "  Method name is: \n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['input_ids'] tensor_info:\n",
      "        dtype: DT_INT32\n",
      "        shape: (-1, 5)\n",
      "        name: serving_default_input_ids:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['output_1'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (-1, 5)\n",
      "        name: StatefulPartitionedCall:0\n",
      "  Method name is: tensorflow/serving/predict\n",
      "\n",
      "Defined Functions:\n",
      "  Function Name: '__call__'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #2\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #3\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #4\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "\n",
      "  Function Name: '_default_save_signature'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}\n",
      "\n",
      "  Function Name: 'call_and_return_all_conditional_losses'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #2\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #3\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n",
      "    Option #4\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}\n",
      "        Named Argument #1\n",
      "          training\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --all --dir ./model/tensorflow/saved_model/0/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-02 01:16:17.976719: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1\n",
      "2021-01-02 01:16:19.289618: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1\n",
      "2021-01-02 01:16:19.368877: E tensorflow/stream_executor/cuda/cuda_driver.cc:314] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected\n",
      "2021-01-02 01:16:19.368919: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (ip-172-16-38-126): /proc/driver/nvidia/version does not exist\n",
      "2021-01-02 01:16:19.369413: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2021-01-02 01:16:19.376557: I tensorflow/core/platform/profile_utils/cpu_utils.cc:104] CPU Frequency: 2999995000 Hz\n",
      "2021-01-02 01:16:19.376889: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x564ac1fc6700 initialized for platform Host (this does not guarantee that XLA will be used). Devices:\n",
      "2021-01-02 01:16:19.376908: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version\n",
      "WARNING:tensorflow:From /home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/tools/saved_model_cli.py:444: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.\n",
      "INFO:tensorflow:Restoring parameters from ./model/tensorflow/saved_model/0/variables/variables\n",
      "2021-01-02 01:16:20.746026: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 93763584 exceeds 10% of free system memory.\n",
      "2021-01-02 01:16:20.746033: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 93763584 exceeds 10% of free system memory.\n",
      "2021-01-02 01:16:20.746027: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 93763584 exceeds 10% of free system memory.\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/bin/saved_model_cli\", line 8, in <module>\n",
      "    sys.exit(main())\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/tools/saved_model_cli.py\", line 1185, in main\n",
      "    args.func(args)\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/tools/saved_model_cli.py\", line 748, in run\n",
      "    init_tpu=args.init_tpu, tf_debug=args.tf_debug)\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/tools/saved_model_cli.py\", line 449, in run_saved_model_with_feed_dict\n",
      "    outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 958, in run\n",
      "    run_metadata_ptr)\n",
      "  File \"/home/ec2-user/anaconda3/envs/python3/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1157, in _run\n",
      "    (np_val.shape, subfeed_t.name, str(subfeed_t.get_shape())))\n",
      "ValueError: Cannot feed value of shape (1, 64) for Tensor 'serving_default_input_ids:0', which has shape '(None, 5)'\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli run --dir ./model/tensorflow/saved_model/0/ --tag_set serve --signature_def serving_default \\\n",
    "    --input_exprs 'input_ids=np.zeros((1,64))' # ;input_mask=np.zeros((1,64))' # ;segment_ids=np.zeros((1,64))'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyze Debugger Rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.latest_training_job.rule_job_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_job_debugger_artifacts_path = estimator.latest_job_debugger_artifacts_path()\n",
    "print(training_job_debugger_artifacts_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyze Debugger Profiling Data\n",
    "\n",
    "Copy outputs of the following cell (`training_job_name` and `region`) to run the analysis notebooks `profiling_generic_dashboard.ipynb`, `analyze_performance_bottlenecks.ipynb`, and `profiling_interactive_analysis.ipynb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_job_name = estimator.latest_training_job.name\n",
    "print(f\"Training jobname: {training_job_name}\")\n",
    "print(f\"Region: {region}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the training is still in progress you can visualize the performance data in SageMaker Studio or in the notebook.\n",
    "Debugger provides utilities to plot system metrics in form of timeline charts or heatmaps. Checkout out the notebook \n",
    "[profiling_interactive_analysis.ipynb](analysis_tools/profiling_interactive_analysis.ipynb) for more details. In the following code cell we plot the total CPU and GPU utilization as timeseries charts. To visualize other metrics such as I/O, memory, network you simply need to extend the list passed to `select_dimension` and `select_events`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smdebug.profiler.analysis.notebook_utils.training_job import TrainingJob\n",
    "tj = TrainingJob(training_job_name, region)\n",
    "tj.wait_for_sys_profiling_data_to_be_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smdebug.profiler.analysis.notebook_utils.timeline_charts import TimelineCharts\n",
    "\n",
    "system_metrics_reader = tj.get_systems_metrics_reader()\n",
    "system_metrics_reader.refresh_event_file_list()\n",
    "\n",
    "view_timeline_charts  = TimelineCharts(system_metrics_reader, \n",
    "                                       framework_metrics_reader=None,\n",
    "                                       select_dimensions=[\"CPU\", \"GPU\"],\n",
    "                                       select_events=[\"total\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Debugger Profiling Report and Notebook\n",
    "\n",
    "The profiling report rule will create an html report `profiler-report.html` with a summary of builtin rules and recommenades of next steps. You can find this report in your S3 bucket.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiler_report = estimator.output_path + estimator.latest_training_job.job_name + \"/rule-output/ProfilerReport/profiler-output/profiler-report.html\"\n",
    "print(profiler_report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp $profiler_report ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiler_notebook = estimator.output_path + estimator.latest_training_job.job_name + \"/rule-output/ProfilerReport/profiler-output/profiler-report.ipynb\"\n",
    "print(profiler_notebook)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp $profiler_notebook ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp s3://$bucket/$training_job_name/output/model.tar.gz ./model.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./model/\n",
    "!tar -xvzf ./model.tar.gz -C ./model/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --all --dir ./model/tensorflow/saved_model/0/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show the Experiment Tracking Lineage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.analytics import ExperimentAnalytics\n",
    "\n",
    "lineage_table = ExperimentAnalytics(\n",
    "    sagemaker_session=sess,\n",
    "    experiment_name=experiment_name,\n",
    "    metric_names=['validation:accuracy'],\n",
    "    sort_by=\"CreationTime\",\n",
    "    sort_order=\"Ascending\",\n",
    ")\n",
    "\n",
    "lineage_df = lineage_table.dataframe()\n",
    "lineage_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lineage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.describe_trial_component(TrialComponentName=lineage_df.TrialComponentName[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pass Variables to the Next Notebook(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store training_job_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store prepare_trial_component_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store training_job_debugger_artifacts_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "Jupyter.notebook.save_checkpoint();\n",
    "Jupyter.notebook.session.delete();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
